import librosa.display
import numpy as np
import matplotlib.pyplot as plt
import pylab
import os
import tqdm

def show_melspectrogram(file_path):
    y, sr = librosa.load(file_path)
    D = np.abs(librosa.stft(y)) ** 2
    S = librosa.feature.melspectrogram(S=D)
    plt.figure(figsize=(10, 4))
    librosa.display.specshow(librosa.power_to_db(S, ref=np.max),
                             y_axis='mel', fmax=8000, x_axis='time')
    plt.colorbar(format='%+2.0f dB')
    plt.title('Mel spectrogram')
    plt.tight_layout()
    plt.show()


def sava_melspectrogram(file_path,save_path):
    sig, fs = librosa.load(file_path)
    pylab.axis('off')
    pylab.axes([0., 0., 1., 1.], frameon=False, xticks=[], yticks=[]) # Remove the white edge
    S = librosa.feature.melspectrogram(y=sig, sr=fs)
    librosa.display.specshow(librosa.power_to_db(S, ref=np.max))
    pylab.savefig(save_path, bbox_inches=None, pad_inches=0)
    pylab.close()

def all_path(dirname):
    result = []
    for maindir, subdir, file_name_list in os.walk(dirname):
        for filename in file_name_list:
            if filename[0]!=".":
                apath = os.path.join(maindir, filename)
                result.append(apath)
    return result

if __name__ == '__main__':
    #os.mkdir("melspectrogram_graph_casia")
    for file in tqdm.tqdm(all_path(("audio_casia"))):
        sava_path=os.path.join("./melspectrogram_graph_casia",file[12:-4])
        #rint(sava_path)
        sava_melspectrogram(file_path=file,save_path=sava_path)